#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Python 2.7 透明隧道客户端
把本机 TCP+UDP 端口截留后，仅通过两条 TCP 长连发到服务器
"""
from __future__ import print_function
import socket, struct, select, threading, subprocess, sys, os

SERVER_IP   = '1.2.3.4'   # 改成你的服务器
TCP_TUNNEL_PORT = 24444
UDP_TUNNEL_PORT = 24445

# ---------- 通用工具 ----------
def forward(sock1, sock2):
    while True:
        r, _, _ = select.select([sock1, sock2], [], [], 60)
        if sock1 in r:
            data = sock1.recv(4096)
            if not data: break
            sock2.sendall(data)
        if sock2 in r:
            data = sock2.recv(4096)
            if not data: break
            sock1.sendall(data)
    sock1.close(); sock2.close()

# ---------- TCP 透明代理 ----------
def tcp_redirect(lport, rport):
    """把本机 lport 全部转到隧道 rport"""
    srv = socket.socket(); srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind(('0.0.0.0', lport)); srv.listen(5)
    while True:
        cli, addr = srv.accept()
        tunnel = socket.create_connection((SERVER_IP, rport))
        # 先写 2 字节长度+2 字节本地端口，方便服务器知道往哪还原
        tunnel.sendall(struct.pack('!HH', lport, 0))  # 0 表示新建流
        threading.Thread(target=forward, args=(cli, tunnel)).start()

# ---------- UDP 透明代理 ----------
def udp_redirect(lport, rport):
    """UDP 截留：每帧前加 4 字节长度+2 字节端口，再塞进 TCP 隧道"""
    udp = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    udp.bind(('0.0.0.0', lport))
    tunnel = socket.create_connection((SERVER_IP, rport))
    def udp2tunnel():
        while True:
            data, addr = udp.recvfrom(65535)
            header = struct.pack('!HH', lport, len(data))
            tunnel.sendall(header + data)
    def tunnel2udp():
        while True:
            header = recv_exact(tunnel, 4)
            port, dlen = struct.unpack('!HH', header)
            data = recv_exact(tunnel, dlen)
            udp.sendto(data, ('127.0.0.1', port))   # 回包给客户端
    threading.Thread(target=udp2tunnel).start()
    threading.Thread(target=tunnel2udp).start()

def recv_exact(sock, n):
    buf = b''
    while len(buf) < n:
        chunk = sock.recv(n - len(buf))
        if not chunk: raise socket.error("EOF")
        buf += chunk
    return buf

# ---------- 平台相关：把 A 进程端口重定向到本机 30000+ 段 ----------
def setup_redirect(ports):
    if sys.platform.startswith('linux'):
        for p in ports:
            subprocess.call(['iptables', '-t', 'nat', '-A', 'OUTPUT',
                             '-p', 'tcp', '--dport', str(p),
                             '-j', 'REDIRECT', '--to-port', str(p+20000)])
            subprocess.call(['iptables', '-t', 'nat', '-A', 'OUTPUT',
                             '-p', 'udp', '--dport', str(p),
                             '-j', 'REDIRECT', '--to-port', str(p+20000)])
    elif sys.platform == 'win32':
        import ctypes, threading, time

        # ---------- WinDivert 1.4 C API 常量 ----------
        WINDIVERT_LAYER_NETWORK = 0
        WINDIVERT_FLAG_SNIFF = 0
        WINDIVERT_FLAG_DROP = 1

        # ---------- 加载 DLL ----------
        try:
            dll = ctypes.WinDLL(r'C:\Users\Lenovo\PyCharmMiscProject\V3\WinDivert.dll')
        except WindowsError:
            raise SystemExit('[!] 请将 WinDivert.dll / WinDivert.sys 放到当前目录')

        # ---------- 函数原型 ----------
        # HANDLE WinDivertOpen(const char *filter, WINDIVERT_LAYER layer, INT16 priority, UINT64 flags)
        dll.WinDivertOpen.argtypes = [ctypes.c_char_p, ctypes.c_int, ctypes.c_short, ctypes.c_uint64]
        dll.WinDivertOpen.restype = ctypes.c_void_p
        # BOOL WinDivertRecv(HANDLE handle, PVOID pPacket, UINT packetLen, PVOID pAddr, PUINT recvLen)
        dll.WinDivertRecv.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_uint, ctypes.c_void_p,
                                      ctypes.POINTER(ctypes.c_uint)]
        dll.WinDivertRecv.restype = ctypes.c_bool
        # BOOL WinDivertSend(HANDLE handle, PVOID pPacket, UINT packetLen, PVOID pAddr, PUINT sendLen)
        dll.WinDivertSend.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_uint, ctypes.c_void_p,
                                      ctypes.POINTER(ctypes.c_uint)]
        dll.WinDivertSend.restype = ctypes.c_bool
        # BOOL WinDivertClose(HANDLE handle)
        dll.WinDivertClose.argtypes = [ctypes.c_void_p]
        dll.WinDivertClose.restype = ctypes.c_bool

        # ---------- 辅助：生成过滤串 ----------
        def build_filter(ports):
            tcp_parts = ['tcp.DstPort == %d' % p for p in ports]
            udp_parts = ['udp.DstPort == %d' % p for p in ports]
            return ' or '.join(tcp_parts + udp_parts)

        # ---------- 后台线程：循环 recv -> 改端口 -> send ----------
        def divert_worker(handle, port_map):
            # 预分配缓冲区
            max_packet = 65535
            packet_buf = ctypes.create_string_buffer(max_packet)
            addr = ctypes.create_string_buffer(32)  # WINDIVERT_ADDRESS 实际 24 字节
            recv_len = ctypes.c_uint()

            while True:
                recv_len.value = max_packet
                if not dll.WinDivertRecv(handle, packet_buf, max_packet, addr, ctypes.byref(recv_len)):
                    time.sleep(0.001)
                    continue
                # 这里只做**最小修改**：把目的端口改为 port+20000
                # IP 头长度 20 字节，TCP/UDP 头偏移 20 字节，目的端口占 2 字节
                if recv_len.value < 22:
                    dll.WinDivertSend(handle, packet_buf, recv_len.value, addr, ctypes.byref(recv_len))
                    continue
                raw = packet_buf.raw[:recv_len.value]
                dst_port = struct.unpack('!H', raw[22:24])[0]
                if dst_port in port_map:
                    new_port = port_map[dst_port]
                    # 写回新端口
                    raw = raw[:22] + struct.pack('!H', new_port) + raw[24:]
                dll.WinDivertSend(handle, raw, len(raw), addr, ctypes.byref(recv_len))

        # ---------- 主逻辑 ----------
        port_map = dict((p, p + 20000) for p in ports)
        filter_str = build_filter(ports)
        handle = dll.WinDivertOpen(filter_str.encode('ascii'),
                                   WINDIVERT_LAYER_NETWORK, 0, 0)
        if not handle:
            raise SystemExit('[!] WinDivertOpen 失败，检查驱动是否已安装或需要管理员权限')
        print('[+] WinDivert 句柄已打开，规则：', filter_str)
        # 启动后台线程
        t = threading.Thread(target=divert_worker, args=(handle, port_map))
        t.daemon = True
        t.start()
    elif sys.platform == 'darwin':
        for p in ports:
            subprocess.call(['echo', 'rdr pass on lo0 inet proto tcp from any to any port %s -> 127.0.0.1 port %s' % (p, p+20000)], stdout=open('/etc/pf.anchors/pytunnel', 'a'))
        subprocess.call(['pfctl', '-f', '/etc/pf.conf'])

# ---------- 主入口 ----------
def main():
    if len(sys.argv) < 2:
        print('Usage: python2 pytunnel.py <port1> [port2] ...')
        sys.exit(1)
    ports = map(int, sys.argv[1:])
    setup_redirect(ports)
    # 为每个原始端口启动 30000+ 的透明代理
    for p in ports:
        threading.Thread(target=tcp_redirect, args=(p+20000, TCP_TUNNEL_PORT)).start()
        threading.Thread(target=udp_redirect, args=(p+20000, UDP_TUNNEL_PORT)).start()
    print('[*] 隧道已启动，Ctrl+C 退出')
    # python3
    while True:
        input('')
    # python2
    # while True: raw_input('')

if __name__ == '__main__':
    main()